#!/bin/bash
# Exit immediately if a command exits with a non-zero status.
set -e

# Move one-level up to tensorflow/models/research directory.
cd ..

# Update PYTHONPATH.
export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim

# Set up the working environment.
CURRENT_DIR=$(pwd)
WORK_DIR="${CURRENT_DIR}/deeplab"

echo "run model_test..."
# Run model_test first to make sure the PYTHONPATH is correctly set.
# python "${WORK_DIR}"/model_test.py -v
echo "model_test successfully..."

# Go to datasets folder and download PASCAL VOC 2012 segmentation dataset.
DATASET_DIR="datasets"
cd "${WORK_DIR}/${DATASET_DIR}"
# sh download_and_convert_bottle.sh

# Go back to original directory.
cd "${CURRENT_DIR}"

# Set up the working directories.
BOTTLE_FOLDER="dummy_bottle"
EXP_FOLDER="exp/train_on_trainval_set"
INIT_FOLDER="${WORK_DIR}/${DATASET_DIR}/${BOTTLE_FOLDER}/init_models"
TRAIN_LOGDIR="${WORK_DIR}/${DATASET_DIR}/${BOTTLE_FOLDER}/${EXP_FOLDER}/train"
EVAL_LOGDIR="${WORK_DIR}/${DATASET_DIR}/${BOTTLE_FOLDER}/${EXP_FOLDER}/eval"
VIS_LOGDIR="${WORK_DIR}/${DATASET_DIR}/${BOTTLE_FOLDER}/${EXP_FOLDER}/vis"
EXPORT_DIR="${WORK_DIR}/${DATASET_DIR}/${BOTTLE_FOLDER}/${EXP_FOLDER}/export"
mkdir -p "${INIT_FOLDER}"
mkdir -p "${TRAIN_LOGDIR}"
mkdir -p "${EVAL_LOGDIR}"
mkdir -p "${VIS_LOGDIR}"
mkdir -p "${EXPORT_DIR}"

# Copy locally the trained checkpoint as the initial checkpoint.
TF_INIT_ROOT="http://download.tensorflow.org/models"
TF_INIT_CKPT="deeplabv3_pascal_train_aug_2018_01_04.tar.gz"
cd "${INIT_FOLDER}"
# wget -nd -c "${TF_INIT_ROOT}/${TF_INIT_CKPT}"
tar -xf "${TF_INIT_CKPT}"
cd "${CURRENT_DIR}"

BOTTLE_DATASET="${WORK_DIR}/${DATASET_DIR}/${BOTTLE_FOLDER}/tfrecord"
echo "start training bottles..."
# Train 1000 iterations.
NUM_ITERATIONS=10000
#python "${WORK_DIR}"/train.py \
#  --logtostderr \
#  --train_split="train" \
#  --model_variant="xception_65" \
#  --atrous_rates=6 \
#  --atrous_rates=12 \
#  --atrous_rates=18 \
#  --output_stride=16 \
#  --decoder_output_stride=4 \
#  --train_crop_size=513 \
#  --train_crop_size=513 \
#  --train_batch_size=4 \
#  --training_number_of_steps="${NUM_ITERATIONS}" \
#  --fine_tune_batch_norm=true \
#  # --tf_initial_checkpoint="${INIT_FOLDER}/deeplabv3_bottle_train_aug/model.ckpt" \
#  --train_logdir="${TRAIN_LOGDIR}" \
#  --dataset_dir="${BOTTLE_DATASET}" \
#  --dataset="bottle"

--logtostderr --train_split="train" --model_variant="xception_65" --atrous_rates=6 --atrous_rates=12 --atrous_rates=18  --output_stride=16 --decoder_output_stride=4  --train_crop_size=513 --train_crop_size=513 --train_batch_size=4 --training_number_of_steps="10000" --fine_tune_batch_norm=false --tf_initial_checkpoint="./datasets/dummy_bottle/init_models/deeplabv3_bottle_train_aug/model.ckpt" --train_logdir="./datasets/dummy_bottle/exp/train" --dataset_dir="./datasets/dummy_bottle/tfrecord" --dataset="bottle"

echo "finish training..."

# Run evaluation. This performs eval over the full val split (1449 images) and
# will take a while.
# Using the provided checkpoint, one should expect mIOU=82.20%.
python "${WORK_DIR}"/eval.py \
  --logtostderr \
  --eval_split="val" \
  --model_variant="xception_65" \
  --atrous_rates=6 \
  --atrous_rates=12 \
  --atrous_rates=18 \
  --output_stride=16 \
  --decoder_output_stride=4 \
  --eval_crop_size=513 \
  --eval_crop_size=513 \
  --checkpoint_dir="${TRAIN_LOGDIR}" \
  --eval_logdir="${EVAL_LOGDIR}" \
  --dataset_dir="${BOTTLE_DATASET}" \
  --max_number_of_evaluations=1

--logtostderr --eval_split="val" --model_variant="xception_65" --atrous_rates=6 --atrous_rates=12 --atrous_rates=18 --output_stride=16  --decoder_output_stride=4  --eval_crop_size=513  --eval_crop_size=513 --checkpoint_dir="./datasets/dummy_bottle/exp/train" --eval_logdir="./datasets/dummy_bottle/exp/eval" --dataset_dir="./datasets/dummy_bottle/tfrecord" --max_number_of_evaluations=1 --dataset=bottle

  --logtostderr --vis_split="val" --model_variant="xception_65" --atrous_rates=6 --atrous_rates=12 --atrous_rates=18 --output_stride=16 --decoder_output_stride=4 --vis_crop_size=513 --vis_crop_size=513 --checkpoint_dir="./datasets/dummy_bottle/exp/train" --vis_logdir="./datasets/dummy_bottle/exp/vis" --dataset_dir="./datasets/dummy_bottle/tfrecord" --max_number_of_iterations=1 --dataset=bottle